/*-
 * Copyright (C) 2019, Shanghai Trusted Industry Control Platform Co., Ltd.
 * Created by 周正辉 (Zhenghui Zhou) <zhouzhenghui@ticpsh.com>
 */

#include "IncrementalParser.h"

using namespace clang;

namespace rocket {
namespace incparser {

IncrementalParser::IncrementalParser(CompilerInstance &CI_)
    : CI(CI_), parser(CI.getPreprocessor(), CI.getSema(), 0) {
  Preprocessor &PP = CI.getPreprocessor();
  //  assert(PP.getCurrentLexer());
  assert(PP.isIncrementalProcessingEnabled());
  parser.Initialize();
}

IncrementalParserContinuation IncrementalParser::parse(
    const std::string &buffer_name, const std::string &buffer,
    clang::DiagnosticConsumer &diag) {
  SourceManager &SM = CI.getSourceManager();
#if LLVM_VERSION_MAJOR < 18
  const clang::FileEntry *FE =
      SM.getFileManager().getVirtualFile(buffer_name, buffer_name.length(), 0);
  SM.overrideFileContents(
      FE, llvm::MemoryBuffer::getMemBufferCopy(buffer, buffer_name));
  SourceLocation include_loc = SM.getLocForEndOfFile(SM.getMainFileID());
  FileID FID = SM.createFileID(FE, include_loc, SrcMgr::C_User);
#else
  auto FFR = SM.getFileManager().getVirtualFileRef(buffer_name,
                                                   buffer_name.length(), 0);
  SM.overrideFileContents(
      FFR, llvm::MemoryBuffer::getMemBufferCopy(buffer, buffer_name));
  SourceLocation include_loc = SM.getLocForEndOfFile(SM.getMainFileID());
  FileID FID = SM.createFileID(FFR, include_loc, SrcMgr::C_User);
#endif
  Preprocessor &PP = CI.getPreprocessor();
  PP.EnterSourceFile(FID, 0, include_loc);

  DiagnosticsEngine &D = CI.getDiagnostics();
  std::unique_ptr<DiagnosticConsumer> orig_diag = D.takeClient();
  D.setClient(&diag, false);

  return IncrementalParserContinuation(*this, std::move(orig_diag));
}

IncrementalParserContinuation IncrementalParser::parse(
    const std::string &code, DiagnosticConsumer &diag) {
  SourceManager &SM = CI.getSourceManager();
  std::string tmp = "_ROCKET_CITMP_FILE";
#if LLVM_VERSION_MAJOR < 18
  const clang::FileEntry *FE =
      SM.getFileManager().getVirtualFile(tmp, tmp.length(), 0);
  SM.overrideFileContents(FE, llvm::MemoryBuffer::getMemBufferCopy(code, tmp));
  SourceLocation include_loc = SM.getLocForEndOfFile(SM.getMainFileID());
  FileID FID = SM.createFileID(FE, include_loc, SrcMgr::C_User);
#else
  auto FFR = SM.getFileManager().getVirtualFileRef(tmp, tmp.length(), 0);
  SM.overrideFileContents(FFR, llvm::MemoryBuffer::getMemBufferCopy(code, tmp));
  SourceLocation include_loc = SM.getLocForEndOfFile(SM.getMainFileID());
  FileID FID = SM.createFileID(FFR, include_loc, SrcMgr::C_User);
#endif
  Preprocessor &PP = CI.getPreprocessor();
  PP.EnterSourceFile(FID, 0, include_loc);

  DiagnosticsEngine &D = CI.getDiagnostics();
  std::unique_ptr<DiagnosticConsumer> orig_diag = D.takeClient();
  D.setClient(&diag, false);

  // TextDiagnosticBuffer *diag_buf = dynamic_cast<TextDiagnosticBuffer
  // *>(D.getClient()); assert(diag_buf != NULL); diag_buf->clear();
  // DiagnosticErrorTrap Trap(D);
  return IncrementalParserContinuation(*this, std::move(orig_diag));
}

IncrementalParserContinuation::~IncrementalParserContinuation() {
  clang::Parser::DeclGroupPtrTy ADecl;
  while (!parser.ParseTopLevelDecl(ADecl)) {
    if (handle) {
      if (!handle(ADecl.get())) handle = nullptr;
    }
  }
  if (error_handle && diag.getNumErrors() > NumErrors) {
    error_handle();
  }

  Preprocessor &PP = parser.CI.getPreprocessor();
#ifdef LLVM_ON_WIN32
  // Microsoft-specific:
  // Late parsed templates can leave unswallowed "macro"-like tokens.
  // They will seriously confuse the Parser when entering the next
  // source file. So lex until we are EOF.
  Token Tok;
  Tok.setKind(tok::eof);
  do {
    PP.Lex(Tok);
  } while (Tok.isNot(tok::eof));
#endif

#ifndef NDEBUG
  Token AssertTok;
  PP.Lex(AssertTok);
  assert(AssertTok.is(tok::eof) &&
         "Lexer must be EOF when starting incremental parse!");
#endif

  parser.CI.getDiagnostics().setClient(orig_diag, true);
}

}  // namespace incparser
}  // namespace rocket
